{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "11fd5648-0757-4050-9ab3-558b7a16712e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import List, Union\n",
    "\n",
    "import pandas as pd\n",
    "from langchain.chains.graph_qa.cypher_utils import CypherQueryCorrector, Schema\n",
    "from langchain_community.graphs import Neo4jGraph\n",
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    SystemMessage,\n",
    "    ToolMessage,\n",
    ")\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    ")\n",
    "from langchain_core.pydantic_v1 import BaseModel\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_google_vertexai import ChatVertexAI\n",
    "\n",
    "# LLMs\n",
    "os.environ[\"GOOGLE_APPLICATION_CREDENTIALS\"] = \"credentials.json\"\n",
    "llm = ChatVertexAI(model_name=\"gemini-1.5-pro\")\n",
    "\n",
    "system = \"\"\"Given an input question, convert it to a Cypher query.\n",
    "To translate a question into a Cypher query, please follow these steps:\n",
    "\n",
    "1. Carefully analyze the provided graph schema to understand what nodes, relationships, and properties are available. Pay attention to the node labels, relationship types, and property keys.\n",
    "2. Identify the key entities and relationships mentioned in the natural language question. Map these to the corresponding node labels, relationship types, and properties in the graph schema.\n",
    "3. Think through how to construct a Cypher query to retrieve the requested information step-by-step. Focus on:\n",
    "   - Identifying the starting node(s) \n",
    "   - Traversing the necessary relationships\n",
    "   - Filtering based on property values\n",
    "   - Returning the requested information\n",
    "Feel free to use multiple MATCH, WHERE, and RETURN clauses as needed.\n",
    "4. Explain how your Cypher query will retrieve the necessary information from the graph to answer the original question. Provide this explanation inside <explanation> tags.\n",
    "5. Once you have finished explaining, construct the Cypher query inside triple backticks ```cypher```.\n",
    "\n",
    "Remember, the goal is to construct a Cypher query that will retrieve the relevant information to answer the question based on the given graph schema.\n",
    "Carefully map the entities and relationships in the question to the nodes, relationships, and properties in the schema.\n",
    "Additional instructions:\n",
    "1. **Array Length**: Always use `size(array)` instead of `length(array)` to get the number of elements in an array.\n",
    "2. **Implicit aggregations**: Always use intermediate WITH clause when performing aggregations\n",
    "3. **Target Neo4j version is 5**: Use Cypher syntax for Neo4j version 5 and above. Do not use any deprecated syntax.\n",
    "\"\"\"\n",
    "\n",
    "# Generate Cypher statement based on natural language input\n",
    "cypher_template = \"\"\"Based on the Neo4j graph schema below, write a Cypher query that would answer the user's question:\n",
    "{schema}\n",
    "\n",
    "Question: {question}\"\"\"  # noqa: E501\n",
    "\n",
    "cypher_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",system ),\n",
    "        (\"human\", cypher_template),\n",
    "    ]\n",
    ")\n",
    "\n",
    "cypher_response = (\n",
    "    cypher_prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "afe825f7-94a3-41f1-8842-a21c26837aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = pd.read_csv(\"gemini_questions.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3385c13-f35d-458a-836a-00c32a70ee72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>type</th>\n",
       "      <th>database</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Which movies from France have an IMDb rating h...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What are the titles of the top 5 highest-gross...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>List the first 3 movies directed by Harold Lloyd.</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What genres are associated with the movie \"Toy...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What is the average IMDb rating of movies rele...</td>\n",
       "      <td>Simple Retrieval Queries</td>\n",
       "      <td>recommendations</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  Which movies from France have an IMDb rating h...   \n",
       "1  What are the titles of the top 5 highest-gross...   \n",
       "2  List the first 3 movies directed by Harold Lloyd.   \n",
       "3  What genres are associated with the movie \"Toy...   \n",
       "4  What is the average IMDb rating of movies rele...   \n",
       "\n",
       "                       type         database  \n",
       "0  Simple Retrieval Queries  recommendations  \n",
       "1  Simple Retrieval Queries  recommendations  \n",
       "2  Simple Retrieval Queries  recommendations  \n",
       "3  Simple Retrieval Queries  recommendations  \n",
       "4  Simple Retrieval Queries  recommendations  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "questions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "62f1720c-6020-4414-89d6-00a05a38c2a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemas = pd.read_csv('text2cypher_schemas.csv')\n",
    "schemas.head()\n",
    "schema_dict = {}\n",
    "for i, row in schemas.iterrows():\n",
    "    schema_dict[row['database']] = row['schema']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "56152468-2572-4f81-997a-df1f70393e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def extract_cypher(text):\n",
    "    # Adjust pattern to capture after ```cypher and spans multiple lines until ```\n",
    "    pattern = r\"```cypher\\n(.*?)\\n```\"\n",
    "    match = re.search(pattern, text, re.DOTALL)\n",
    "\n",
    "    if match:\n",
    "        # Return the extracted text if triple backticks are present\n",
    "        return match.group(1).strip()\n",
    "    else:\n",
    "        # Return the original text if triple backticks are not present\n",
    "        return None\n",
    "\n",
    "def extract_explanation(text):\n",
    "    pattern = re.compile(r'<explanation>(.*?)</explanation>', re.DOTALL)\n",
    "    match = pattern.search(text)\n",
    "    if match:\n",
    "        explanation_content = match.group(1).strip()\n",
    "        return explanation_content\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "35dc3988-987c-4f20-8d73-e4dffea565fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "cypher_responses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c1092b-5039-4b54-9db3-12c1aac31d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, row in questions.iterrows():\n",
    "    if i % 50 == 0:\n",
    "        print(i)\n",
    "    schema = schema_dict[row[\"database\"]]\n",
    "    try:\n",
    "        output = cypher_response.invoke({\"question\": row[\"question\"], \"schema\": schema})\n",
    "        cypher_responses.append(\n",
    "            {\n",
    "                \"question\": row[\"question\"],\n",
    "                \"database\": row[\"database\"],\n",
    "                \"output\": output,\n",
    "                \"type\": row[\"type\"],\n",
    "                \"cypher\": extract_cypher(output),\n",
    "                \"explanation\": extract_explanation(output)\n",
    "            }\n",
    "        )\n",
    "    except Exception as e:\n",
    "        output = cypher_response.invoke({\"question\": row[\"question\"], \"schema\": schema})\n",
    "        cypher_responses.append(\n",
    "            {\n",
    "                \"question\": row[\"question\"],\n",
    "                \"database\": row[\"database\"],\n",
    "                \"output\": output,\n",
    "                \"type\": row[\"type\"],\n",
    "                \"cypher\": extract_cypher(output),\n",
    "                \"explanation\": extract_explanation(output)\n",
    "            }\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fd2cac-df69-454f-807f-db71981ec689",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = pd.DataFrame.from_records(cypher_responses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf2d32f-9e37-40ef-b0fc-61c4767cdf3b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "DEMO_URL = \"neo4j+s://demo.neo4jlabs.com\"\n",
    "\n",
    "def create_graph(database):\n",
    "    return Neo4jGraph(\n",
    "        url=DEMO_URL,\n",
    "        username=database,\n",
    "        password=database,\n",
    "        database=database,\n",
    "        refresh_schema=False,\n",
    "        timeout=10,\n",
    "    )\n",
    "\n",
    "syntax_error = []\n",
    "returns_results = []\n",
    "timeouts = []\n",
    "not_possible = []\n",
    "last_graph = \"\"\n",
    "for i, row in results.iterrows():\n",
    "    if i % 100 == 0:\n",
    "        print(i)\n",
    "    \n",
    "    # To avoid a new driver for every request\n",
    "    if row[\"database\"] != last_graph:\n",
    "        last_graph = row[\"database\"]\n",
    "        print(last_graph)\n",
    "        graph = create_graph(row[\"database\"])\n",
    "    if not isinstance(row['cypher'],str) or row[\"cypher\"].startswith(\"//\"):\n",
    "            returns_results.append(False)\n",
    "            syntax_error.append(False)\n",
    "            timeouts.append(False)\n",
    "            not_possible.append(True)\n",
    "    else:\n",
    "        not_possible.append(False)\n",
    "        try:\n",
    "            data = graph.query(row[\"cypher\"])\n",
    "            if data:\n",
    "                returns_results.append(True)\n",
    "            else:\n",
    "                returns_results.append(False)\n",
    "            syntax_error.append(False)\n",
    "            timeouts.append(False)\n",
    "        except ValueError as e:\n",
    "            if \"Generated Cypher Statement is not valid\" in str(e):\n",
    "                syntax_error.append(True)\n",
    "                print(f\"Syntax error in Cypher query: {e}\")\n",
    "            else:\n",
    "                syntax_error.append(False)\n",
    "                print(f\"Other ValueError: {e}\")\n",
    "            returns_results.append(False)\n",
    "            timeouts.append(False)\n",
    "        except Exception as e:\n",
    "            if (\n",
    "                hasattr(e, 'code') and e.code\n",
    "                == \"Neo.ClientError.Transaction.TransactionTimedOutClientConfiguration\"\n",
    "            ):\n",
    "                returns_results.append(False)\n",
    "                syntax_error.append(False)\n",
    "                timeouts.append(True)\n",
    "            else:\n",
    "                returns_results.append(False)\n",
    "                syntax_error.append(False)\n",
    "                timeouts.append(True)\n",
    "                # Some weird errors we create a new graph object\n",
    "                try:\n",
    "                    graph._driver.close()\n",
    "                except:\n",
    "                    pass\n",
    "                graph = create_graph(row[\"database\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c109c7f2-3f2f-494a-a7cd-5f611852c3a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "output[\"syntax_error\"] = syntax_error\n",
    "output[\"timeout\"] = timeouts\n",
    "output[\"returns_results\"] = returns_results\n",
    "output[\"no_cypher\"] = not_possible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c990ad76-4b77-461d-9414-44675cff286d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distribution for syntax_error:\n",
      " syntax_error\n",
      "False    8769\n",
      "True      876\n",
      "Name: count, dtype: int64\n",
      "Distribution for timeout:\n",
      " timeout\n",
      "False    9330\n",
      "True      315\n",
      "Name: count, dtype: int64\n",
      "Distribution for returns_results:\n",
      " returns_results\n",
      "True     6658\n",
      "False    2987\n",
      "Name: count, dtype: int64\n",
      "Distribution for no_cypher:\n",
      " no_cypher\n",
      "False    9441\n",
      "True      204\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Assume df is your DataFrame and col1, col2, col3 are the boolean columns\n",
    "distribution_col1 = output[\"syntax_error\"].value_counts()\n",
    "distribution_col2 = output[\"timeout\"].value_counts()\n",
    "distribution_col3 = output[\"returns_results\"].value_counts()\n",
    "distribution_col4 = output[\"no_cypher\"].value_counts()\n",
    "\n",
    "\n",
    "print(\"Distribution for syntax_error:\\n\", distribution_col1)\n",
    "print(\"Distribution for timeout:\\n\", distribution_col2)\n",
    "print(\"Distribution for returns_results:\\n\", distribution_col3)\n",
    "print(\"Distribution for no_cypher:\\n\", distribution_col4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "54c66866-b90a-4719-81ce-e15e97522329",
   "metadata": {},
   "outputs": [],
   "source": [
    "output.to_csv('text2cypher_gemini.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f7e083-8390-49b1-acf2-fef7343b5e79",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
